"""
Kraken Upscale & Tile Calc
MIT License © 2025 The Kraken (@KrakenUnbound)
"""
import math, os, re, time, logging
import numpy as np, torch
from PIL import Image, ImageDraw
import folder_paths

counter = 0

def _round_to(x, base): return int(round(x / base) * base)
def _ceil_to(x, base): return int(math.ceil(x / base) * base)
def _clamp(v, lo, hi): return max(lo, min(hi, v))
def _blank_image_tensor(w=8, h=8): return torch.zeros((1, h, w, 3), dtype=torch.float32)

class KrakenUpscaleTileCalc:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {"resolution_in": ("STRING", {"default": "1024 x 1024"})},
            "optional": {
                "preset_in": ("STRING", {"default": "Custom"}),
                "image": ("IMAGE",),
                "scale_step": ("FLOAT", {"default": 0.01, "min": 0.005, "max": 0.25, "step": 0.005}),
                "crop_margin_px": ("INT", {"default": 32, "min": 0, "max": 128, "step": 1}),
                "mask_blur": ("INT", {"default": 16, "min": 0, "max": 64, "step": 1}),
                "tile_padding": ("INT", {"default": 16, "min": 0, "max": 256, "step": 1}),
                "seam_fix_width": ("INT", {"default": 64, "min": 0, "max": 256, "step": 1}),
                "seam_fix_mask_blur": ("INT", {"default": 16, "min": 0, "max": 64, "step": 1}),
                "seam_fix_padding": ("INT", {"default": 32, "min": 0, "max": 256, "step": 1}),
                "min_upscale_policy": (["auto","off","1.5x","2.0x","custom"], {"default": "auto"}),
                "min_upscale_custom": ("FLOAT", {"default": 1.5, "min": 1.0, "max": 4.0, "step": 0.1}),
            }
        }

    RETURN_TYPES = ("IMAGE","FLOAT","INT","INT","INT","INT","INT","INT","INT","STRING","STRING","STRING")
    RETURN_NAMES  = ("image_out","upscale_by","tile_width","tile_height","mask_blur","tile_padding","seam_fix_width","seam_fix_mask_blur","seam_fix_padding","predicted_resolution","target_resolution","target_preset")
    FUNCTION = "calculate"
    CATEGORY = "kraken/upscale"

    def _parse_resolution_ws(self, s):
        s = (s or "").lower().replace(" ", "")
        if "x" in s:
            try:
                w, h = s.split("x"); return (int(w), int(h))
            except Exception: return None
        return None

    def _parse_resolution_from_label(self, label):
        s = str(label or ""); m = re.search(r"(\d+)\s*[xX]\s*(\d+)", s)
        return (int(m.group(1)), int(m.group(2))) if m else None

    def _image_size(self, image, resolution_in):
        if image is not None:
            if isinstance(image, torch.Tensor):
                if image.ndim == 4: _, h, w, _ = image.shape
                else: h, w = image.shape[0], image.shape[1]
                return (w, h)
            if isinstance(image, np.ndarray):
                if image.ndim == 4: _, h, w, _ = image.shape
                else: h, w = image.shape[0], image.shape[1]
                return (w, h)
        return self._parse_resolution_ws(resolution_in) or (1024,1024)

    def _choose_min_upscale(self, wi, hi, Wt, Ht, policy, custom):
        policy = (policy or "auto").lower()
        if policy == "off": return 1.0
        if policy == "1.5x": return 1.5
        if policy == "2.0x": return 2.0
        if policy == "custom":
            try: return max(1.0, min(4.0, float(custom)))
            except Exception: return 1.5
        src_mp = (wi*hi)/1_000_000.0; tgt_mp = (Wt*Ht)/1_000_000.0
        if max(Wt,Ht) <= 1024 or tgt_mp <= 1.0: return 2.0
        if src_mp <= 2.0: return 2.0
        return 1.5

    def _make_preview(self, pred_w, pred_h, Wt, Ht):
        w, h = max(pred_w, Wt), max(pred_h, Ht)
        scale = max(1, min(512 / w, 512 / h))
        canvas = Image.new("RGB", (int(w*scale), int(h*scale)), (24,24,24))
        draw = ImageDraw.Draw(canvas)
        draw.rectangle([0,0,canvas.size[0]-1, canvas.size[1]-1], outline=(160,160,160), width=2)
        tw, th = int(Wt*scale), int(Ht*scale)
        x0, y0 = (canvas.size[0]-tw)//2, (canvas.size[1]-th)//2
        draw.rectangle([x0,y0,x0+tw-1,y0+th-1], outline=(0,200,120), width=2)
        draw.text((6,6), f"pred {pred_w}x{pred_h}", fill=(220,220,220))
        draw.text((6,24), f"target {Wt}x{Ht}", fill=(0,200,120))
        return canvas

    def calculate(self, resolution_in, preset_in="Custom", image=None, scale_step=0.01, crop_margin_px=32,
                  mask_blur=16, tile_padding=16, seam_fix_width=64, seam_fix_mask_blur=16, seam_fix_padding=32,
                  min_upscale_policy="auto", min_upscale_custom=1.5):

        wi, hi = self._image_size(image, resolution_in)
        tgt = self._parse_resolution_ws(preset_in) or self._parse_resolution_from_label(preset_in) or self._parse_resolution_ws(resolution_in) or (wi, hi)
        Wt, Ht = tgt

        crop_margin_px = int(crop_margin_px); crop_margin_px = _clamp(crop_margin_px, 0, max(Wt,Ht)//10)
        Wm, Hm = Wt + 2*crop_margin_px, Ht + 2*crop_margin_px

        required = max(Wm/wi, Hm/hi)
        s = max(1.0, round((required/scale_step)) * scale_step)

        min_s = self._choose_min_upscale(wi, hi, Wt, Ht, min_upscale_policy, min_upscale_custom)
        if min_s > 1.0:
            s = max(s, round((min_s/scale_step)) * scale_step)
        s = round(float(s), 3)

        pred_align = 8
        pred_w = _ceil_to(wi * s, pred_align)
        pred_h = _ceil_to(hi * s, pred_align)

        bump_guard = 0
        while (pred_w < Wt or pred_h < Ht) and bump_guard < 4:
            s = round(s + scale_step, 3)
            pred_w = _ceil_to(wi * s, pred_align)
            pred_h = _ceil_to(hi * s, pred_align)
            bump_guard += 1

        overlap = _clamp(int(round(min(pred_w, pred_h) * 0.015)), 32, 96)
        tile_w = _round_to(_clamp(pred_w // 2 + overlap, 512, 2048), 64)
        tile_h = _round_to(_clamp(pred_h // 2 + overlap, 512, 2048), 64)

        predicted_resolution = f"{pred_w} x {pred_h}"
        target_resolution = f"{Wt} x {Ht}"
        target_preset = f"{Wt}x{Ht}"

        global counter; counter += 1
        from folder_paths import get_temp_directory
        file = f"kraken_tilecalc_{time.time()}_{counter:05}.png"
        try:
            preview_img = self._make_preview(pred_w, pred_h, Wt, Ht)
            preview_img.save(os.path.join(get_temp_directory(), file), compress_level=4)
            images = [{"filename": file, "subfolder": "", "type": "temp"}]
        except Exception as e:
            logging.warning(f"Preview generation failed: {e}"); images = []

        image_out = image if image is not None else _blank_image_tensor()
        return {"ui": {"images": images},
                "result": (image_out, s, tile_w, tile_h, mask_blur, tile_padding, seam_fix_width, seam_fix_mask_blur,
                            seam_fix_padding, predicted_resolution, target_resolution, target_preset)}
